#include "main/utils.h"
#include "lib/binary_tree.h"
using namespace std;

int depthOfBinaryTree(BinaryTreeNode* root) {
  if (!root)
    return 0;
  int depth = max(depthOfBinaryTree(root->left), depthOfBinaryTree(root->right)) + 1;
  return depth;
}

bool isBalancedCore(BinaryTreeNode* root, int &depth) {
  if (!root)
    return true;
  
  int depth_left = 0, depth_right = 0;
  bool left = isBalancedCore(root->left, depth_left);
  bool right = isBalancedCore(root->right, depth_right);
  int diff = depth_right - depth_left;
  if (left && right && diff >= -1 && diff <= 1) {
    depth = max(depth_left, depth_right) + 1;
    return true;
  }
  return false;
}

bool isBalanced(BinaryTreeNode* root) {
  int depth = 0;
  return isBalancedCore(root, depth);
}

int main() {
  vector<int> vec = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
  auto root = CreateBinaryTree(vec);
  int depth = depthOfBinaryTree(root);
  cout << "The depth of binary tree is: " << depth << endl;

  bool balance = isBalanced(root);
  cout << "The binary tree is balanced: " << balance << endl;

  BinaryTreeNode* node1 = CreateBinaryTreeNode(1);
  BinaryTreeNode* node2 = CreateBinaryTreeNode(2);
  BinaryTreeNode* node3 = CreateBinaryTreeNode(3);
  BinaryTreeNode* node4 = CreateBinaryTreeNode(4);
  BinaryTreeNode* node5 = CreateBinaryTreeNode(5);
  BinaryTreeNode* node6 = CreateBinaryTreeNode(6);

  ConnectNode(node1, node2, node3);
  ConnectNode(node2, node4, nullptr);
  ConnectNode(node4, node5, nullptr);
  ConnectNode(node3, nullptr, node6);
  bool no_balance = isBalanced(node1);
  cout << "The binary tree is not balanced: " << no_balance << endl;


  return 0;
}
